/*****************************************************************************
 * process.c : process command one by one
 *****************************************************************************
 * Copyright (C) 2011 PELinker
 *
 * Authors: Iceberg <iceberg58744@163.com>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston MA 02110-1301, USA.
 *****************************************************************************/

#include "common.h"
#include "util.h"

#define EXE "EXE"
#define DLL "DLL"

#define NATIVE "native"
#define WINDOWS "windows"
#define CONSOLE "console"

#define BASE_ALIGN_FACTOR (64 * 1024)
#define MIN_ALIGN_FACTOR 512
#define MAX_ALIGN_FACTOR (64 * 1024)
#define MAX_VERSION 255

static const String FORMAT_EXE =
{
	sizeof(EXE) - 1,
	EXE
};

static const String FORMAT_DLL =
{
	sizeof(DLL) - 1,
	DLL
};

static const String SUBSYSTEM_NATIVE =
{
	sizeof(NATIVE) - 1,
	NATIVE
};

static const String SUBSYSTEM_WINDOWS =
{
	sizeof(WINDOWS) - 1,
	WINDOWS
};

static const String SUBSYSTEM_CONSOLE =
{
	sizeof(CONSOLE) - 1,
	CONSOLE
};

#define RegisterOptionHandler(opt_name, op) \
{ \
	{ \
		sizeof(opt_name) - 1, \
		(opt_name) \
	}, \
	(op) \
}

typedef int (*OptionProcessor)(const Option * const opt, Link_Option * const link_opt);

typedef struct _OptionHandler
{
	String name;
	OptionProcessor proc;
} OptionHandler;

static int ConvertHexStringToDWORD(const String * const src, DWORD *dest)
{
	DWORD i = 0;

	*dest = 0;

	while (i < src->length)
	{
		if (src->ptr[i] >= '0' && src->ptr[i] <= '9')
		{
			*dest <<= 4;
			*dest += (src->ptr[i] - '0');
		}
		else if (src->ptr[i] >= 'A' && src->ptr[i] <= 'F')
		{
			*dest <<= 4;
			*dest += (src->ptr[i] - 'A' + 10);
		}
		else if (src->ptr[i] >= 'a' && src->ptr[i] <= 'f')
		{
			*dest <<= 4;
			*dest += (src->ptr[i] - 'a' + 10);
		}
		else
		{
			return ERR_INVALID_ARGUMENT;
		}

		++i;
	}

	return 0;
}

static int ConvertDecimalStringToWORDS(const String * const src, WORD *dest1, WORD *dest2)
{
	DWORD i = 0;

	WORD *dest = dest1;

	*dest1 = 0;

	*dest2 = 0;

	while (i < src->length)
	{
		if (src->ptr[i] >= '0' && src->ptr[i] <= '9')
		{
			*dest *= 10;
			*dest += src->ptr[i] - '0';
		}
		else if ('.' == src->ptr[i])
		{
			dest = dest2;
		}
		else
		{
			return ERR_INVALID_ARGUMENT;
		}

		++i;
	}

	return 0;
}

static int SetPrintHelpInfo(const Option * const opt, Link_Option * const link_opt)
{
	link_opt->bPrintHelpInfo = 1;

	return 0;
}

static int SetMapFileName(const Option * const opt, Link_Option * const link_opt)
{
	if (link_opt->bSetMapFileName)
	{
		return ERR_DUPLICATED_OPTION;
	}

	if (!opt->value.length)
	{
		return ERR_INCOMPLETE_OPTION;
	}

	link_opt->MapFileName = opt->value;

	link_opt->bSetMapFileName = 1;

	return 0;
}

static int SetOutputFileName(const Option * const opt, Link_Option * const link_opt)
{
	if (link_opt->bSetOutputFileName)
	{
		return ERR_DUPLICATED_OPTION;
	}

	if (!opt->value.length)
	{
		return ERR_INCOMPLETE_OPTION;
	}

	link_opt->OutputFileName = opt->value;

	link_opt->bSetOutputFileName = 1;

	return 0;
}

static int SetOutputFileType(const Option * const opt, Link_Option * const link_opt)
{
	if (link_opt->bSetOutputFileType)
	{
		return ERR_DUPLICATED_OPTION;
	}

	if (!opt->value.length)
	{
		return ERR_INCOMPLETE_OPTION;
	}

	if (!StringCompare(&(opt->value), &FORMAT_EXE))
	{
		link_opt->OutputFileType = IMAGE_FILE_EXECUTABLE_IMAGE;
	}
	else if (!StringCompare(&(opt->value), &FORMAT_DLL))
	{
		link_opt->OutputFileType = IMAGE_FILE_DLL;
	}
	else
	{
		return ERR_INVALID_ARGUMENT;
	}

	link_opt->bSetOutputFileType = 1;

	return 0;
}

static int SetEntryPoint(const Option * const opt, Link_Option * const link_opt)
{
	if (link_opt->bSetEntryPoint)
	{
		return ERR_DUPLICATED_OPTION;
	}

	if (!opt->value.length)
	{
		return ERR_INCOMPLETE_OPTION;
	}

	link_opt->EntryPoint = opt->value;

	link_opt->bSetEntryPoint = 1;

	return 0;
}

static int SetImageBase(const Option * const opt, Link_Option * const link_opt)
{
	if (link_opt->bSetImageBase)
	{
		return ERR_DUPLICATED_OPTION;
	}

	if (!opt->value.length)
	{
		return ERR_INCOMPLETE_OPTION;
	}

	if (ConvertHexStringToDWORD(&(opt->value), &(link_opt->ImageBase)))
	{
		return ERR_INVALID_ARGUMENT;
	}

	if (!link_opt->ImageBase || (link_opt->ImageBase & (BASE_ALIGN_FACTOR - 1)))
	{
		return ERR_INVALID_ARGUMENT;
	}

	link_opt->bSetImageBase = 1;

	return 0;
}

static int SetSectionAlign(const Option * const opt, Link_Option * const link_opt)
{
	if (link_opt->bSetSectionAlign)
	{
		return ERR_DUPLICATED_OPTION;
	}

	if (!opt->value.length)
	{
		return ERR_INCOMPLETE_OPTION;
	}

	if (ConvertHexStringToDWORD(&(opt->value), &(link_opt->SectionAlign)))
	{
		return ERR_INVALID_ARGUMENT;
	}

	if (!link_opt->SectionAlign || (link_opt->SectionAlign & (link_opt->SectionAlign - 1)))
	{
		return ERR_INVALID_ARGUMENT;
	}

	link_opt->bSetSectionAlign = 1;

	return 0;
}

static int SetFileAlign(const Option * const opt, Link_Option * const link_opt)
{
	if (link_opt->bSetFileAlign)
	{
		return ERR_DUPLICATED_OPTION;
	}

	if (!opt->value.length)
	{
		return ERR_INCOMPLETE_OPTION;
	}

	if (ConvertHexStringToDWORD(&(opt->value), &(link_opt->FileAlign)))
	{
		return ERR_INVALID_ARGUMENT;
	}

	if (!link_opt->FileAlign || (link_opt->FileAlign & (link_opt->FileAlign - 1)) || link_opt->FileAlign < MIN_ALIGN_FACTOR || link_opt->FileAlign > MAX_ALIGN_FACTOR)
	{
		return ERR_INVALID_ARGUMENT;
	}

	link_opt->bSetFileAlign = 1;

	return 0;
}

static int SetSubsystem(const Option * const opt, Link_Option * const link_opt)
{
	if (link_opt->bSetSubsystem)
	{
		return ERR_DUPLICATED_OPTION;
	}

	if (!opt->value.length)
	{
		return ERR_INCOMPLETE_OPTION;
	}

	if (!StringCompare(&(opt->value), &SUBSYSTEM_NATIVE))
	{
		link_opt->Subsystem = IMAGE_SUBSYSTEM_NATIVE;
	}
	else if (!StringCompare(&(opt->value), &SUBSYSTEM_WINDOWS))
	{
		link_opt->Subsystem = IMAGE_SUBSYSTEM_WINDOWS_GUI;
	}
	else if (!StringCompare(&(opt->value), &SUBSYSTEM_CONSOLE))
	{
		link_opt->Subsystem = IMAGE_SUBSYSTEM_WINDOWS_CUI;
	}
	else
	{
		return ERR_INVALID_ARGUMENT;
	}

	link_opt->bSetSubsystem = 1;

	return 0;
}

static int SetLinkerVersion(const Option * const opt, Link_Option * const link_opt)
{
	if (link_opt->bSetLinkerVersion)
	{
		return ERR_DUPLICATED_OPTION;
	}

	if (!opt->value.length)
	{
		return ERR_INCOMPLETE_OPTION;
	}

	if (ConvertDecimalStringToWORDS(&(opt->value), &(link_opt->LinkerVersion[0]), &(link_opt->LinkerVersion[1])))
	{
		return ERR_INVALID_ARGUMENT;
	}

	if (link_opt->LinkerVersion[0] > MAX_VERSION || link_opt->LinkerVersion[1] > MAX_VERSION)
	{
		return ERR_INVALID_ARGUMENT;
	}

	link_opt->bSetLinkerVersion = 1;

	return 0;
}

static int SetOSVersion(const Option * const opt, Link_Option * const link_opt)
{
	if (link_opt->bSetOSVersion)
	{
		return ERR_DUPLICATED_OPTION;
	}

	if (!opt->value.length)
	{
		return ERR_INCOMPLETE_OPTION;
	}

	if (ConvertDecimalStringToWORDS(&(opt->value), &(link_opt->OSVersion[0]), &(link_opt->OSVersion[1])))
	{
		return ERR_INVALID_ARGUMENT;
	}

	if (link_opt->OSVersion[0] > MAX_VERSION || link_opt->OSVersion[1] > MAX_VERSION)
	{
		return ERR_INVALID_ARGUMENT;
	}

	link_opt->bSetOSVersion = 1;

	return 0;
}

static int SetImageVersion(const Option * const opt, Link_Option * const link_opt)
{
	if (link_opt->bSetImageVersion)
	{
		return ERR_DUPLICATED_OPTION;
	}

	if (!opt->value.length)
	{
		return ERR_INCOMPLETE_OPTION;
	}

	if (ConvertDecimalStringToWORDS(&(opt->value), &(link_opt->ImageVersion[0]), &(link_opt->ImageVersion[1])))
	{
		return ERR_INVALID_ARGUMENT;
	}

	if (link_opt->ImageVersion[0] > MAX_VERSION || link_opt->ImageVersion[1] > MAX_VERSION)
	{
		return ERR_INVALID_ARGUMENT;
	}

	link_opt->bSetImageVersion = 1;

	return 0;
}

static int SetSubSysVersion(const Option * const opt, Link_Option * const link_opt)
{
	if (link_opt->bSetSubSysVersion)
	{
		return ERR_DUPLICATED_OPTION;
	}

	if (!opt->value.length)
	{
		return ERR_INCOMPLETE_OPTION;
	}

	if (ConvertDecimalStringToWORDS(&(opt->value), &(link_opt->SubSysVersion[0]), &(link_opt->SubSysVersion[1])))
	{
		return ERR_INVALID_ARGUMENT;
	}

	if (link_opt->SubSysVersion[0] > MAX_VERSION || link_opt->SubSysVersion[1] > MAX_VERSION)
	{
		return ERR_INVALID_ARGUMENT;
	}

	link_opt->bSetSubSysVersion = 1;

	return 0;
}

static int SetStackSize(const Option * const opt, Link_Option * const link_opt)
{
	if (link_opt->bSetStackSize)
	{
		return ERR_DUPLICATED_OPTION;
	}

	if (!opt->value.length)
	{
		return ERR_INCOMPLETE_OPTION;
	}

	if (ConvertHexStringToDWORD(&(opt->value), &(link_opt->StackSize)))
	{
		return ERR_INVALID_ARGUMENT;
	}

	link_opt->bSetStackSize = 1;

	return 0;
}

static int SetStackCommitSize(const Option * const opt, Link_Option * const link_opt)
{
	if (link_opt->bSetStackCommitSize)
	{
		return ERR_DUPLICATED_OPTION;
	}

	if (!opt->value.length)
	{
		return ERR_INCOMPLETE_OPTION;
	}

	if (ConvertHexStringToDWORD(&(opt->value), &(link_opt->StackCommitSize)))
	{
		return ERR_INVALID_ARGUMENT;
	}

	link_opt->bSetStackCommitSize = 1;

	return 0;
}

static int SetHeapSize(const Option * const opt, Link_Option * const link_opt)
{
	if (link_opt->bSetHeapSize)
	{
		return ERR_DUPLICATED_OPTION;
	}

	if (!opt->value.length)
	{
		return ERR_INCOMPLETE_OPTION;
	}

	if (ConvertHexStringToDWORD(&(opt->value), &(link_opt->HeapSize)))
	{
		return ERR_INVALID_ARGUMENT;
	}

	link_opt->bSetHeapSize = 1;

	return 0;
}

static int SetHeapCommitSize(const Option * const opt, Link_Option * const link_opt)
{
	if (link_opt->bSetHeapCommitSize)
	{
		return ERR_DUPLICATED_OPTION;
	}

	if (!opt->value.length)
	{
		return ERR_INCOMPLETE_OPTION;
	}

	if (ConvertHexStringToDWORD(&(opt->value), &(link_opt->HeapCommitSize)))
	{
		return ERR_INVALID_ARGUMENT;
	}

	link_opt->bSetHeapCommitSize = 1;

	return 0;
}

static const OptionHandler OptionHandlerArray[] =
{
	RegisterOptionHandler("h", SetPrintHelpInfo),
	RegisterOptionHandler("H", SetPrintHelpInfo),
	RegisterOptionHandler("?", SetPrintHelpInfo),
	RegisterOptionHandler("mapfile", SetMapFileName),
	RegisterOptionHandler("output", SetOutputFileName),
	RegisterOptionHandler("format", SetOutputFileType),
	RegisterOptionHandler("entry", SetEntryPoint),
	RegisterOptionHandler("base", SetImageBase),
	RegisterOptionHandler("memalign", SetSectionAlign),
	RegisterOptionHandler("filealign", SetFileAlign),
	RegisterOptionHandler("subsystem", SetSubsystem),
	RegisterOptionHandler("linkerver", SetLinkerVersion),
	RegisterOptionHandler("osver", SetOSVersion),
	RegisterOptionHandler("imagever", SetImageVersion),
	RegisterOptionHandler("subsysver", SetSubSysVersion),
	RegisterOptionHandler("stacksize", SetStackSize),
	RegisterOptionHandler("stackcommitsize", SetStackCommitSize),
	RegisterOptionHandler("heapsize", SetHeapSize),
	RegisterOptionHandler("heapcommitsize", SetHeapCommitSize)
};

static int ProcessOption(const Option * const opt, Link_Option * const link_opt)
{
	int i;

	for (i = 0; i < sizeof(OptionHandlerArray) / sizeof(OptionHandlerArray[0]); ++i)
	{
		if (!StringCompare(&OptionHandlerArray[i].name, &(opt->name)))
		{
			if (OptionHandlerArray[i].proc)
			{
				return (OptionHandlerArray[i].proc)(opt, link_opt);
			}
		}
	}

	return ERR_UNRECOGNIZED_OPTION;
}

int ProcessCommand(const Command * const cmd, Link_Option * const link_opt)
{
	DWORD i;

	if (!cmd->TotalOptionNumber)
	{
		return SetPrintHelpInfo(NULL, link_opt);
	}

	for (i = 0; i < cmd->TotalOptionNumber; ++i)
	{
		int ret = ProcessOption(cmd->options + i, link_opt);

		if (ret)
		{
			if (ERR_DUPLICATED_OPTION == ret)
			{
				PrintInfoAux("Duplicated option: ");
			}
			else
			{
				PrintInfoAux("Invalid option: ");
			}
			PrintInfoEx(cmd->options[i].name.ptr, cmd->options[i].name.length);
			if (cmd->options[i].value.ptr)
			{
				PrintInfoAux(":");
				PrintInfoEx(cmd->options[i].value.ptr, cmd->options[i].value.length);
			}
			PrintInfoAux("\n");

			return ret;
		}
	}

	return 0;
}
